##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(parallel)
library(dplyr)
library(parallel)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--geneset_file"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ_peak-gene/testis_combined_peak.combineRNA.qc.Rdata"

    geneset_file <- "~/20231121_singleMuti/results/celltype_plot/trajectory/positive/pct_0.25.list"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/trajectory/positive/pct_0.25/enrichment"

}


###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
out_path <- opt$out_path
geneset_file <- opt$geneset_file

dir.create(out_path , recursive = T)
image_path <- out_path

###########################################################################################

a <- load(comine_data_file)
## testis_combined_peak_combineRNA
projHeme5 <- testis_combined_peak_combineRNA

dat_geneset <- data.frame(fread(geneset_file , header = F))

cpu <- 20

###########################################################################################
## 所有peak
all_peak <- data.frame(projHeme5@peakSet)
all_peak$cell_type <- sapply( strsplit(all_peak$GroupReplicate , "._.") , "[" , 1 )
all_peak$peakName <- paste0( all_peak$seqnames , ":" , all_peak$start , "-" , all_peak$end )

###########################################################################################
## 注释motif所在peak
motifPositions <- getMatches(projHeme5, name="Motif")

all_motif_match <- Reduce(function(x,y)bind_rows(x,y),mclapply( colnames(motifPositions) ,function(motif){

    print(motif)
    match_pos <- which(assay(motifPositions)[,motif])
    tmp_peak <- all_peak[match_pos,]
    tmp_peak$motif_in <- motif

    return(tmp_peak)
},mc.cores=cpu))


image_name <- paste0( out_path , "/motif_in_peak-gene.tsv")
write.table( all_motif_match , image_name , row.names = F , sep = "\t" , quote = F )

###########################################################################################
## 提取160个关注的motif
all_motif_match$gene <- sapply(strsplit(all_motif_match$motif_in , "_") , "[" , 1)
all_motif_match_use <- subset( all_motif_match , gene %in% dat_geneset$V1 )

## 对motif，两两计算是否共同出现在同一peak上面

result <- Reduce(function(x,y)bind_rows(x,y),mclapply( 1:(length(dat_geneset$V1)-1) ,function(i){

    #print(i)
    motif1 <- dat_geneset$V1[i]
    print(motif1)

    tmp_res <- c()
    for( j in (i+1):length(dat_geneset$V1)){
        motif2 <- dat_geneset$V1[j]
        print(motif2)

        tmp_peak <- subset(all_motif_match_use , gene %in% c(motif1 , motif2))
        tmp_peak <- unique(tmp_peak[,c("peakName" , "gene")])

        ## 共有的peak
        share_peak <- names(which(table(tmp_peak$peakName) == 2))

        ## 两种共有的peak
        a <- length(share_peak)
        ## motif1独有的
        b <- nrow(subset( tmp_peak , gene == motif1 & !(peakName %in% share_peak) ))
        ## motif2独有的
        c <- nrow(subset( tmp_peak , gene == motif2 & !(peakName %in% share_peak) ))
        ## 其它所有peak
        d <- length(unique(all_motif_match$peakName)[!unique(all_motif_match$peakName) %in% tmp_peak$peakName])
        ## 矩阵
        data <- matrix( c(a,b,c,d) , ncol = 2 )

        ## 卡方检验
        chisq_result <- chisq.test(data)
        p_value <- chisq_result$p.value

        ## 富集系数
        enrichment <- data[1,1] / (data[1,1] + data[1,2]) / (data[1,1] + data[2,1]) * (data[1,1] + data[2,1] + data[1,2] + data[2,2])

        tmp_dat <- data.frame( motif1 = motif1 , motif2 = motif2 , conPeakNum = a , motif1_uniq = b , motif2_uniq = c , other = d , P = p_value , ER = enrichment )
        tmp_res <- rbind( tmp_res , tmp_dat )
    }
    return(tmp_res)

},mc.cores=cpu))


image_name <- paste0( out_path , "/importantTF_enrichment.tsv")
write.table( result , image_name , row.names = F , sep = "\t" , quote = F )